//
// @author 风杨
// @date ${DATE} ${TIME}
// @see <a href="yueny09@163.com">mailTo: yueny09@163.com</a>
//
use std::fs;
use datafusion::prelude::*;
use datafusion::dataframe::*;
use parquet::file::reader::{FileReader, SerializedFileReader};
use parquet::record::RowAccessor;

#[global_allocator]
static ALLOC: snmalloc_rs::SnMalloc = snmalloc_rs::SnMalloc;

#[tokio::main]
async fn main() -> datafusion::error::Result<()> {
    // 带列头的数据文件
    // head is : id, seq, money, bri, level, desc, flag, memo
    let csv_path = "order_table_with_head.csv";
    println!("{:?}", csv_path);

    let ctx = SessionContext::new();

    // create the dataframe
    let mut options = CsvReadOptions::new();
    options = options.has_header(false);
    ctx.register_csv("order", csv_path, options).await?;

    let sql;
    {
        sql = "select * from order \
    union all select * from order \
    union all select * from order ";
    }

    // create a plan to run a SQL query
    let df = ctx.sql(sql).await?;
    // execute and print results
    df.show().await?;

    let record_batch_vec = df.collect().await?;
    let count: usize = record_batch_vec
        .iter()
        .map(|v| v.num_rows())
        .sum();
    println!("查询总条数： {}", count);


    println!("开始输出 parquet 文件");
    let _ = fs::remove_dir_all("parquet");

    // TODO 写入时应指定 WriterProperties， 暂时不可行。报错
    // let write_properties: Option<WriterProperties> = WriterProperties::builder()
    //     .set_write_batch_size(8192)
    //     .build();

    // <parquet> 为文件夹
    // df.write_parquet("parquet", Some(write_properties)).await?;
    df.write_parquet(".parquet", None).await?;
    println!("parquet 文件写入成功");


    Ok(())
}
